Transformers LogitsProcessor

LogitsProcessor 是一个处理生成模型的输出 logits 的类，它在 Hugging Face Transformers 库中用于概率生成模型的概率分布的后处理，例如 GPT-2 或 BERT。通过使用不同的 logits 处理器，我们可以对模型的输出进行各种修改和调整。

导入库和模块

from transformers import GPT2LMHeadModel, GPT2Tokenizer
from transformers.generation.logits_process import LogitsProcessor, TemperatureLogitsWarper, TopPLogitsWarper

创建模型和 tokenizer

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

创建 LogitsProcessor

processor = LogitsProcessor()

此处我们创建了一个空的 LogitsProcessor，它不会对 logits 进行任何修改。然而，实际上，你将经常使用子类来实现特定的处理行为。

添加处理器

LogitsProcessor 可以包含多个处理器。让我们添加一些常用的处理器：

temperature_processor = TemperatureLogitsWarper(temperature=0.7)
top_p_processor = TopPLogitsWarper(top_p=0.9)

processor.append(temperature_processor)
processor.append(top_p_processor)

这里，我们添加了两个处理器：一个用于修改 logits 的温度，另一个用于实现 Top-P 采样。

使用处理器

现在，我们可以在生成过程中使用 LogitsProcessor：

input_ids = tokenizer.encode("Once upon a time", return_tensors="pt")

output_sequences = model.generate(
    input_ids=input_ids,
    do_sample=True,
    max_length=100,
    pad_token_id=tokenizer.eos_token_id,
    logits_processor=processor,
)

print(tokenizer.decode(output_sequences[0]))

这段代码会生成一个包含 100 个 token 的文本，该文本基于 "Once upon a time" 这个前缀。在生成过程中，我们的 logits 处理器会对模型的输出进行处理。

注意：生成模型的具体行为会取决于具体的 LogitsProcessor 以及它们的参数设置。

本文作者：Maeiee

本文链接：Transformers LogitsProcessor

版权声明：如无特别声明，本文即为原创文章，版权归 Maeiee 所有，未经允许不得转载！

喜欢我文章的朋友请随缘打赏，鼓励我创作更多更好的作品！